# 完成对数据的加载,加载到内存中
import glob
import random  # 随机抽取若干个元素（无重复）

import torch
from torch.utils.data import DataLoader, Dataset
import torchaudio


class MyDataset(Dataset):
    def __init__(self, data_list):
        super(MyDataset, self).__init__()
        # 构建数据x,y
        self.xs = []
        self.ys = []
        for path in data_list:
            x = torchaudio.load(path)[0].reshape(-1, 1)
            # print(path)
            # print(path.split('/')[-1][0])
            y = int(path.split('\\')[-1][0])
            self.xs.append(x)
            self.ys.append(y)
        self.ys = torch.tensor(self.ys)

    def __getitem__(self, index):
        return self.xs[index], self.ys[index]

    def __len__(self):
        return len(self.xs)


datas_path = set(glob.glob('../data/*.wav'))

train_data_path = set(random.sample(list(datas_path), int(len(datas_path) * 0.8)))
test_data_path = datas_path - train_data_path

train_dataset = MyDataset(list(train_data_path))
test_dataset = MyDataset(list(test_data_path))

train_data_loader = DataLoader(dataset=train_dataset, batch_size=8, shuffle=True, num_workers=4, drop_last=True)
test_data_loader = DataLoader(dataset=test_dataset, batch_size=8, shuffle=False, num_workers=4)

if __name__ == '__main__':
    x, y = next(iter(train_data_loader))
    print(x.shape)
    print(y)
